import argparse  # 导入 argparse 模块，用于解析命令行参数
import queue  # 导入 queue 模块，用于实现线程安全的队列
import threading  # 导入 threading 模块，用于多线程操作
import time  # 导入 time 模块，用于延时操作
import sounddevice as sd  # 导入 sounddevice 库，用于音频输入
import numpy as np  # 导入 NumPy 库，用于处理音频数据
import torch  # 导入 PyTorch 库，用于深度学习模型推理
from crnn import WakeWordCRNN  # 导入自定义的 CRNN 模型（WakeWordCRNN）
from audio_processor import AudioProcessor  # 导入自定义的音频预处理工具


class RealTimeDetector:  # 定义实时检测器类
    def __init__(self, args):  # 初始化方法，接收命令行参数
        self.device = torch.device(args.device)  # 根据传入的参数选择设备（CPU 或 GPU）

        # 加载模型
        self.model = WakeWordCRNN().to(self.device)  # 初始化 CRNN 模型，并将其移动到指定设备
        self.model.load_state_dict(torch.load(args.model_path, map_location=self.device))  # 加载预训练模型权重
        self.model.eval()  # 设置模型为评估模式

        # 音频参数
        self.sr = 16000  # 设置采样率为 16kHz
        self.window = 3  # 检测窗口长度（秒），表示每次检测使用 3 秒音频
        self.stride = 1  # 检测间隔（秒），表示每隔 1 秒进行一次检测

        # 环形缓冲区
        self.buffer = np.zeros(int(self.window * self.sr), dtype=np.float32)  # 创建一个长度为 window 的环形缓冲区
        self.q = queue.Queue()  # 创建线程安全的队列，用于存储音频片段
        self.stop_event = threading.Event()  # 创建停止事件，用于控制线程退出

        # 检测参数
        self.threshold = args.threshold  # 设置触发阈值，默认为 0.85
        self.min_trigger = 2  # 连续触发次数，表示连续两次检测概率超过阈值时触发唤醒

    def audio_callback(self, indata, frames, time, status):
        """音频输入回调"""
        self.buffer = np.roll(self.buffer, -frames)  # 将缓冲区向前滚动 frames 个样本
        self.buffer[-frames:] = indata[:, 0]  # 将新输入的音频数据填充到缓冲区末尾
        self.q.put(self.buffer.copy())  # 将当前缓冲区内容复制并放入队列中

    def predict_worker(self):
        """预测线程"""
        trigger_count = 0  # 初始化触发计数器
        while not self.stop_event.is_set():  # 循环运行，直到收到停止事件
            try:
                audio = self.q.get(timeout=1)  # 从队列中获取音频数据，超时时间为 1 秒
                # 预处理
                features = AudioProcessor.extract_features(
                    audio, self.sr, max_len=100
                )  # 提取音频特征
                tensor = torch.FloatTensor(features).unsqueeze(0).to(self.device)  # 转换为张量并移动到指定设备

                # 推理
                with torch.no_grad():  # 禁用梯度计算，节省内存
                    output = self.model(tensor)  # 前向传播，计算模型输出
                    prob = torch.softmax(output, dim=1)[0][1].item()  # 获取唤醒词的概率

                # 触发逻辑
                if prob > self.threshold:  # 如果概率超过阈值
                    trigger_count += 1  # 增加触发计数
                    if trigger_count >= self.min_trigger:  # 如果连续触发次数达到要求
                        print(f"\033[92m! 主人我在！ ({prob:.1%})\033[0m")  # 打印唤醒信息
                        trigger_count = 0  # 重置触发计数
                else:
                    trigger_count = max(0, trigger_count - 1)  # 否则减少触发计数

                print(f"Current score: {prob:.1%}", end='\r')  # 实时显示当前概率

            except queue.Empty:  # 如果队列为空，继续循环
                continue

    def start(self):
        """启动检测"""
        print(f"Starting detection (threshold={self.threshold})...")  # 打印启动信息
        try:
            # 启动预测线程
            predict_thread = threading.Thread(target=self.predict_worker)  # 创建预测线程
            predict_thread.start()  # 启动预测线程

            # 启动音频流
            with sd.InputStream(
                    samplerate=self.sr,  # 设置采样率
                    channels=1,  # 单声道输入
                    callback=self.audio_callback,  # 设置音频回调函数
                    blocksize=int(self.stride * self.sr)  # 设置每次读取的样本数
            ):
                while True:  # 主线程保持运行
                    time.sleep(0.1)  # 每隔 0.1 秒检查一次
        except KeyboardInterrupt:  # 捕获键盘中断信号
            self.stop_event.set()  # 设置停止事件
            predict_thread.join()  # 等待预测线程退出
            print("\nStopped.")  # 打印停止信息


if __name__ == "__main__":  # 主程序入口
    parser = argparse.ArgumentParser()  # 创建命令行参数解析器
    parser.add_argument("--model_path", required=True)  # 添加模型路径参数
    parser.add_argument("--threshold", type=float, default=0.85)  # 添加触发阈值参数
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")  # 添加设备参数
    args = parser.parse_args()  # 解析命令行参数

    detector = RealTimeDetector(args)  # 创建实时检测器实例
    detector.start()  # 启动检测
    #这段代码实现了一个实时语音唤醒词检测系统，能够从麦克风捕获音频并实时判断是否检测到唤醒词。它结合了音频处理、深度学习推理和多线程技术，提供了一个高效且易用的解决方案。